"""
strategy_selector.py定义了StrategySelector类，
通过调用大语言模型根据用户查询选择最合适的检索策略。支持的策略包括直接检索、HyDE（假设问题检索）、子查询检索和回溯问题检索，旨在优化检索阶段的输入处理。
"""
# 导包
import os, sys
from langchain.prompts import PromptTemplate
from openai import OpenAI

from base.logger import get_logger
from base.config import get_config

class StrategySelector:
    """
    StrategySelector类定义了策略选择逻辑，通过调用大语言模型根据用户查询选择最合适的检索策略。
    支持的策略包括直接检索、HyDE（假设问题检索）、子查询检索和回溯问题检索，旨在优化检索阶段的输入处理。
    """
    def __init__(self):
        """
        初始化策略选择器，加载配置文件和日志记录器。
        """
        self.config = get_config()
        self.logger = get_logger()

        # 初始化OpenAI客户端
        self.client = OpenAI(api_key=self.config.DASHSCOPE_API_KEY,
                             base_url=self.config.DASHSCOPE_BASE_URL)

        # 获取策略选择提示模板
        self.strategy_prompt_template = self._get_strategy_prompt()

    def _get_strategy_prompt(self):
        """
        获取策略选择提示模板，用于大语言模型调用。
        """
        #   定义私有方法，获取策略选择 Prompt 模板
        return PromptTemplate(
            template="""
                    你是一个智能助手，负责分析用户查询 {query}，并从以下四种检索增强策略中选择一个最适合的策略，直接返回策略名称，不需要解释过程。

                    以下是几种检索增强策略及其适用场景：

                    1.  **直接检索：**
                        * 描述：对用户查询直接进行检索，不进行任何增强处理。
                        * 适用场景：适用于查询意图明确，需要从知识库中检索**特定信息**的问题，例如：
                            * 示例：
                                * 查询：AI 学科学费是多少？
                                * 策略：直接检索
                            * 查询：JAVA的课程大纲是什么？
                                * 策略：直接检索
                    2.  **假设问题检索（HyDE）：**
                        * 描述：使用 LLM 生成一个假设的答案，然后基于假设答案进行检索。
                        * 适用场景：适用于查询较为抽象，直接检索效果不佳的问题，例如：
                            * 示例：
                                * 查询：人工智能在教育领域的应用有哪些？
                                * 策略：假设问题检索
                    3.  **子查询检索：**
                        * 描述：将复杂的用户查询拆分为多个简单的子查询，分别检索并合并结果。
                        * 适用场景：适用于查询涉及多个实体或方面，需要分别检索不同信息的问题，例如：
                            * 示例：
                                * 查询：比较 Milvus 和 Zilliz Cloud 的优缺点。
                                * 策略：子查询检索
                    4.  **回溯问题检索：**
                        * 描述：将复杂的用户查询转化为更基础、更易于检索的问题，然后进行检索。
                        * 适用场景：适用于查询较为复杂，需要简化后才能有效检索的问题，例如：
                            * 示例：
                                * 查询：我有一个包含 100 亿条记录的数据集，想把它存储到 Milvus 中进行查询。可以吗？
                                * 策略：回溯问题检索

                    根据用户查询 {query}，直接返回最适合的策略名称，例如 "直接检索"。不要输出任何分析过程或其他内容。
                    """
            ,
            input_variables=["query"],
        )
    def call_dashscope(self, prompt: str):
        """
        调用大语言模型，根据输入提示生成文本。
        """
        # 调用大语言模型API生成对话响应
        # 该函数向指定的LLM服务发送聊天请求，包含系统角色设定和用户提示词
        response = self.client.chat.completions.create(
            model=self.config.LLM_MODEL,
            messages=[
                {"role": "system", "content": "你是一个有用的助手。"},
                {"role": "user", "content": prompt}],
            temperature=0.1,
        )
        # 从模型响应中提取生成的文本
        # 假设响应包含一个choices列表，每个元素是一个包含message的字典，message包含role和content
        print("response->", response)
        print("---------------------------")
        print("response.choices->", response.choices)
        print("---------------------------")
        print("response.choices[0]->", response.choices[0])
        print("---------------------------")
        print("response.choices[0].message->", response.choices[0].message)
        generated_text = response.choices[0].message.content.strip()
        return generated_text

    # 选择检索策略
    def select_strategy(self, query: str):
        """
        选择检索策略，根据用户查询调用大语言模型进行策略选择。
        """
        result = self.call_dashscope(self.strategy_prompt_template.format(query=query))
        self.logger.info(f"为查询{query}选择的策略是：{result}")

if __name__ == '__main__':
    strategy_selector = StrategySelector()
    strategy_selector.select_strategy('你好吗')
    # strategy_selector.select_strategy("Mysql数据库能不能支持100w个样本的插入?")